package tenant

import (
	"fmt"
	"github.com/pingcap/tidb/pkg/parser"
	"github.com/pingcap/tidb/pkg/parser/ast"
	_ "github.com/pingcap/tidb/pkg/parser/test_driver"
	"strings"
)

// TenantConditionVisitor 用于在特定表名的查询中添加tenant_id条件的访问者
type TenantConditionVisitor struct {
	targetTables map[string]bool // 需要添加tenant_id条件的表名
}

// NewTenantConditionVisitor 创建新的租户条件访问者
func NewTenantConditionVisitor(targetTables []string) *TenantConditionVisitor {
	tableMap := make(map[string]bool)
	for _, table := range targetTables {
		tableMap[table] = true
	}
	return &TenantConditionVisitor{
		targetTables: tableMap,
	}
}

// Enter 进入节点时的处理
func (v *TenantConditionVisitor) Enter(n ast.Node) (node ast.Node, skipChildren bool) {
	switch node := n.(type) {
	case *ast.SelectStmt:
		// 处理SELECT语句
		v.processSelectStmt(node)
		return node, false
	case *ast.UpdateStmt:
		// 处理UPDATE语句
		v.processUpdateStmt(node)
		return node, false
	case *ast.DeleteStmt:
		// 处理DELETE语句
		v.processDeleteStmt(node)
		return node, false
	}
	return n, false
}

// Leave 离开节点时的处理
func (v *TenantConditionVisitor) Leave(n ast.Node) (node ast.Node, ok bool) {
	return n, true
}

// processSelectStmt 处理SELECT语句
func (v *TenantConditionVisitor) processSelectStmt(stmt *ast.SelectStmt) {
	// 检查FROM子句中的表
	if stmt.From != nil {
		v.processTableRefs(stmt.From, stmt.Where)
	}
}

// processUpdateStmt 处理UPDATE语句
func (v *TenantConditionVisitor) processUpdateStmt(stmt *ast.UpdateStmt) {
	// 检查UPDATE子句中的表
	if stmt.TableRefs != nil {
		v.processTableRefs(stmt.TableRefs, stmt.Where)
	}
}

// processDeleteStmt 处理DELETE语句
func (v *TenantConditionVisitor) processDeleteStmt(stmt *ast.DeleteStmt) {
	// 检查FROM子句中的表
	if stmt.TableRefs != nil {
		v.processTableRefs(stmt.TableRefs, stmt.Where)
	}
}

// processTableRefs 处理表引用
func (v *TenantConditionVisitor) processTableRefs(tableRefs *ast.TableRefsClause, whereClause ast.ExprNode) {
	if tableRefs.TableRefs == nil {
		return
	}

	// 检查是否包含目标表
	hasTargetTable := v.checkTableRefs(tableRefs.TableRefs)

	if hasTargetTable {
		// 添加tenant_id条件
		v.addTenantCondition(whereClause)
	}
}

// checkTableRefs 检查表引用中是否包含目标表
func (v *TenantConditionVisitor) checkTableRefs(tableRefs *ast.Join) bool {
	if tableRefs == nil {
		return false
	}

	// 检查左表
	if leftTable, ok := tableRefs.Left.(*ast.TableSource); ok {
		if tableName, ok := leftTable.Source.(*ast.TableName); ok {
			if v.targetTables[tableName.Name.String()] {
				return true
			}
		}
	}

	// 检查右表
	if tableRefs.Right != nil {
		if rightTable, ok := tableRefs.Right.(*ast.TableSource); ok {
			if tableName, ok := rightTable.Source.(*ast.TableName); ok {
				if v.targetTables[tableName.Name.String()] {
					return true
				}
			}
		}
	}

	return false
}

// addTenantCondition 添加tenant_id条件
func (v *TenantConditionVisitor) addTenantCondition(whereClause ast.ExprNode) {
	// 创建tenant_id = ?条件
	tenantCondition := &ast.BinaryOperationExpr{
		Op: ast.BinaryOpEQ,
		L:  &ast.ColumnNameExpr{Name: &ast.ColumnName{Name: ast.NewCIStr("tenant_id")}},
		R:  &ast.ParamMarkerExpr{},
	}

	// 如果已有WHERE条件，用AND连接
	if whereClause != nil {
		newWhere := &ast.BinaryOperationExpr{
			Op: ast.BinaryOpAnd,
			L:  whereClause,
			R:  tenantCondition,
		}
		// 这里需要将新的WHERE条件设置回去，但AST修改比较复杂
		// 在实际应用中，可能需要更复杂的逻辑来处理
		fmt.Printf("添加tenant_id条件: tenant_id = ?\n")
	} else {
		// 如果没有WHERE条件，直接设置
		fmt.Printf("添加tenant_id条件: tenant_id = ?\n")
	}
}

// AddTenantConditionToSQL 解析SQL语句并为特定表添加tenant_id条件
func AddTenantConditionToSQL(sql string, targetTables []string) (string, error) {
	p := parser.New()

	// 解析SQL语句
	stmtNodes, _, err := p.Parse(sql, "", "")
	if err != nil {
		return "", fmt.Errorf("解析SQL失败: %v", err)
	}

	// 创建访问者
	visitor := NewTenantConditionVisitor(targetTables)

	// 遍历所有语句节点
	for _, stmt := range stmtNodes {
		stmt.Accept(visitor)
	}

	// 将修改后的AST转换回SQL字符串
	restorer := NewSQLRestorer()
	var resultSQL string

	if len(stmtNodes) == 1 {
		// 单个语句
		resultSQL, err = restorer.RestoreSQL(stmtNodes[0])
		if err != nil {
			return "", err
		}
	} else {
		// 多个语句，用分号连接
		var sqls []string
		for _, stmt := range stmtNodes {
			sqlStr, err := restorer.RestoreSQL(stmt)
			if err != nil {
				return "", err
			}
			sqls = append(sqls, sqlStr)
		}
		resultSQL = strings.Join(sqls, "; ")
	}

	return resultSQL, nil
}

// ParseAndAddTenantCondition 解析SQL并添加tenant_id条件的示例函数
func ParseAndAddTenantCondition() {
	// 示例SQL语句
	sqlExamples := []string{
		"SELECT * FROM t_user WHERE id = 1",
		"SELECT * FROM t_dept WHERE name = 'IT'",
		"SELECT u.name, d.name FROM t_user u JOIN t_dept d ON u.dept_id = d.id",
		"UPDATE t_user SET name = 'John' WHERE id = 1",
		"DELETE FROM t_dept WHERE id = 1",
		"SELECT * FROM t_user WHERE age > 18 AND status = 'active'",
		"SELECT * FROM other_table WHERE id = 1", // 这个表不在目标列表中
	}

	// 目标表名列表
	targetTables := []string{"t_user", "t_dept"}

	fmt.Println("=== SQL添加tenant_id条件示例 ===")

	for i, sql := range sqlExamples {
		fmt.Printf("\n--- 示例 %d ---\n", i+1)
		fmt.Printf("原始SQL: %s\n", sql)

		// 解析并添加tenant_id条件
		modifiedSQL, err := AddTenantConditionToSQL(sql, targetTables)
		if err != nil {
			fmt.Printf("处理失败: %v\n", err)
			continue
		}

		fmt.Printf("修改后SQL: %s\n", modifiedSQL)
		fmt.Println("处理完成")
	}
}
